# feature_extraction.py
import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder
import argparse
import os

# 定义标准氨基酸列表 (20种标准氨基酸 + X 表示未知或非标准)
AMINO_ACIDS = sorted(['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 
                      'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X'])

def load_data(file_path):
    """
    从 CSV 文件加载数据。
    Args:
        file_path (str): CSV文件的路径。
    Returns:
        pandas.DataFrame: 加载的数据。
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"错误: 数据文件 '{file_path}' 未找到。")
    try:
        df = pd.read_csv(file_path)
        print(f"数据从 '{file_path}' 加载成功。")
        print(f"数据形状: {df.shape}")
        print(f"数据列名: {df.columns.tolist()}")
        # 确保必要的列存在
        required_cols = ['Peptide_Sequence', 'First3AA', 'No_DPP4', 'DPP4_Digested']
        for col in required_cols:
            if col not in df.columns:
                raise ValueError(f"错误: 数据文件中缺少必需的列 '{col}'。")
        return df
    except Exception as e:
        raise ValueError(f"加载数据时出错: {e}")

def create_target_variable(df):
    """
    根据 DPP4_Digested 和 No_DPP4 列创建目标变量 'Resistant'。
    如果 DPP4_Digested >= No_DPP4，则 Resistant 为 1 (耐受)，否则为 0 (不耐受)。
    Args:
        df (pandas.DataFrame): 包含 DPP4_Digested 和 No_DPP4 列的数据。
    Returns:
        pandas.DataFrame: 增加了 'Resistant' 列的数据。
    """
    if not all(col in df.columns for col in ['DPP4_Digested', 'No_DPP4']):
        raise ValueError("错误: 'DPP4_Digested' 或 'No_DPP4' 列在DataFrame中不存在。")
    
    # 确保这两列是数值类型
    df['DPP4_Digested'] = pd.to_numeric(df['DPP4_Digested'], errors='coerce')
    df['No_DPP4'] = pd.to_numeric(df['No_DPP4'], errors='coerce')

    # 处理可能的NaN值 (例如，如果转换失败或原始数据中存在)
    df.dropna(subset=['DPP4_Digested', 'No_DPP4'], inplace=True)
    
    df['Resistant'] = (df['DPP4_Digested'] >= df['No_DPP4']).astype(int)
    print("目标变量 'Resistant' 创建成功。")
    print(f"Resistant 分布:\n{df['Resistant'].value_counts(normalize=True)}")
    return df

def one_hot_encode_first3aa(df, amino_acids_list=AMINO_ACIDS, fitted_encoder=None): # Add fitted_encoder argument
    """
    对 First3AA 列进行独热编码。
    如果提供了 fitted_encoder，则使用它进行转换；否则，拟合一个新的编码器。
    Args:
        df (pandas.DataFrame): 包含 First3AA 列的数据。
        amino_acids_list (list): 用于编码的氨基酸列表。
        fitted_encoder (OneHotEncoder, optional): 已拟合的 OneHotEncoder 对象。默认为 None。
    Returns:
        pandas.DataFrame: 包含独热编码特征的新DataFrame。
        OneHotEncoder: 使用的 (拟合的或传入的) sklearn OneHotEncoder对象。
    """
    if 'First3AA' not in df.columns:
        raise ValueError("错误: 'First3AA' 列在DataFrame中不存在。")

    df_copy = df.copy() # 操作副本以避免 SettingWithCopyWarning
    df_copy['First3AA'] = df_copy['First3AA'].astype(str).fillna('XXX')

    aa_cols_for_encoding = []
    for i in range(3):
        col_name = f'AA{i+1}'
        df_copy[col_name] = df_copy['First3AA'].apply(lambda x: x[i] if len(x) == 3 and x[i] in amino_acids_list else 'X')
        aa_cols_for_encoding.append(col_name)
    
    input_for_encoder = df_copy[aa_cols_for_encoding]

    encoder_to_use = fitted_encoder
    if encoder_to_use is None:
        print("未提供 fitted_encoder，将拟合新的 OneHotEncoder。")
        encoder_to_use = OneHotEncoder(categories=[amino_acids_list]*3, sparse_output=False, handle_unknown='ignore')
        encoded_aa = encoder_to_use.fit_transform(input_for_encoder)
    else:
        print("使用提供的 fitted_encoder 进行转换。")
        encoded_aa = encoder_to_use.transform(input_for_encoder)
    
    # 创建编码后的特征名称
    # 确保与编码器的 categories_ 属性或 get_feature_names_out 方法一致
    if hasattr(encoder_to_use, 'get_feature_names_out'):
        try:
            # 假设编码器在训练时是以 'AA1', 'AA2', 'AA3' 作为输入特征名（如果适用）
            # 或者直接使用 categories 来构造
            prefix = [f'AA{j+1}' for j in range(len(encoder_to_use.categories_))]
            encoded_feature_names = encoder_to_use.get_feature_names_out(prefix)
        except TypeError: # get_feature_names_out might not take input_features if fitted on unnamed data
            encoded_feature_names = []
            for i, categories_for_pos in enumerate(encoder_to_use.categories_):
                for category_val in categories_for_pos:
                    encoded_feature_names.append(f'AA{i+1}_{category_val}')
    else: # Fallback for older scikit-learn or simpler construction
        encoded_feature_names = []
        for i, categories_for_pos in enumerate(encoder_to_use.categories_): # encoder.categories_ is a list of arrays
            for category_val in categories_for_pos:
                encoded_feature_names.append(f'AA{i+1}_{category_val}')

    # 确保列名数量与编码输出的特征数量匹配
    if encoded_aa.shape[1] != len(encoded_feature_names):
        print(f"警告: 编码后特征数量 ({encoded_aa.shape[1]}) 与生成的列名数量 ({len(encoded_feature_names)}) 不符。请检查编码器和列名生成逻辑。")
        # 作为一种健壮性措施，如果数量不匹配，截断或填充列名，但这表明存在潜在问题。
        # 对于预测/评估，最关键的是编码器能正确转换。
        encoded_feature_names = encoded_feature_names[:encoded_aa.shape[1]]


    encoded_df = pd.DataFrame(encoded_aa, columns=encoded_feature_names, index=df_copy.index)
    
    if fitted_encoder is None:
        print("First3AA 列独热编码 (新拟合) 完成。")
    else:
        print("First3AA 列独热编码 (使用已拟合编码器) 完成。")
        
    print(f"独热编码后特征数量: {encoded_df.shape[1]}")
    
    return encoded_df, encoder_to_use
    
def preprocess_data(file_path, output_dir="processed_data"):
    """
    完整的数据预处理流程：加载、创建目标变量、独热编码。
    Args:
        file_path (str): 原始数据CSV文件的路径。
        output_dir (str): 保存处理后数据的目录。
    Returns:
        pandas.DataFrame: 包含特征和目标变量的完整DataFrame。
        OneHotEncoder: 用于First3AA编码的OneHotEncoder对象。
    """
    print("开始数据预处理...")
    df = load_data(file_path)
    df = create_target_variable(df)
    
    # 分离特征和目标，以避免在编码器中包含目标列
    if 'Resistant' not in df.columns:
        raise ValueError("目标变量 'Resistant' 未成功创建。")
        
    y = df['Resistant']
    X_other_features = df.drop(columns=['Resistant', 'First3AA', 'No_DPP4', 'DPP4_Digested', 'Peptide_Sequence', 'log2FoldChange'], errors='ignore')


    encoded_aa_df, aa_encoder = one_hot_encode_first3aa(df.copy()) # 使用df的副本以避免SettingWithCopyWarning

    # 合并独热编码特征和其他可能保留的特征 (如果有)
    X_processed = pd.concat([X_other_features.reset_index(drop=True), encoded_aa_df.reset_index(drop=True)], axis=1)
    
    # 将目标变量添加回处理后的数据帧
    final_df = pd.concat([X_processed, y.reset_index(drop=True)], axis=1)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    output_path = os.path.join(output_dir, "preprocessed_data.csv")
    final_df.to_csv(output_path, index=False)
    print(f"预处理后的数据已保存到: {output_path}")
    print(f"最终处理后的数据形状: {final_df.shape}")
    
    return final_df, aa_encoder

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="肽序列特征提取脚本")
    parser.add_argument("--input_file", type=str, required=True, help="输入的CSV数据文件路径 (例如: DPP-IV_test.csv)")
    parser.add_argument("--output_dir", type=str, default="processed_data", help="保存预处理后数据的目录")
    
    args = parser.parse_args()
    
    try:
        processed_df, _ = preprocess_data(args.input_file, args.output_dir)
        print("特征提取成功完成。")
        print("处理后的数据预览 (前5行):")
        print(processed_df.head())
    except Exception as e:
        print(f"特征提取过程中发生错误: {e}")

